﻿using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Text.RegularExpressions;

namespace TfsCloudDeploy.Tasks
{
    public static class PEHeaderUtil
    {

        #region Constants

        private const string PE_HEADER_TOOLPATH = @"C:\Program Files\TfsCloudDeploy\Tools\PEedit.exe";
        private const int MAXIMUM_PROCESS_EXIT_WAIT_MILLISECONDS = 180000; // 3 minutes.
        private const int FAILURE_EXIT_CODE = 101;
        private const string MESSAGE_COMMAND_CURRENTLY_EXECUTED = "Current command: {0} {1}.";
        private const string MESSAGE_FAILED_COMMAND_TIMEOUT = "Timeout expired waiting for process to exit (waited {0} milliseconds).";
        private const string MESSAGE_FAILED_RUNNING_COMMAND_PROCESS = "Result: {0}. Failed running command [Exit Code: {1}]: {2} {3}.";
        private const string REGEX_TIMEDATESTAMP = @"TimeDateStamp \(from file header\)\: (?<TIMEDATESTAMP>\w{8})\.";
        private const string REGEX_MVID_FILE_HEADER = @"Current MVID \(from module table\)\: (?<MVID_FILE_HEADER>[-{}\w]{38})\.";
        private const string REGEX_MVID_DEBUG_DIRECTORY_ENTRY = @"Current MVID \(from Debug Directory entry\)\: (?<MVID_FILE_HEADER>[-{}\w]{38})\.";
        private const string REGEX_NEW_TIMEDATESTAMP = @"New TimeDateStamp\: (?<TIMEDATESTAMP>\w{8})\.";
        private const string REGEX_NEW_MVID_FILE_HEADER = @"New MVID \(file header\)\: (?<MVID_FILE_HEADER>[-{}\w]{38})\.";
        private const string REGEX_NEW_MVID_DEBUG_DIRECTORY_ENTRY = @"New MVID \(Debug Directory\)\: (?<MVID_FILE_HEADER>[-{}\w]{38})\.";
        private const string REGEX_FILE_VERSION = @"1\)\s+FileVersion:\s{5}(?<FILEVERSION>\d+\.\d+\.\d+\.\d+)\.";
        private const string REGEX_PRODUCT_VERSION = @"1\)\s+ProductVersion:\s{2}(?<PRODUCTVERSION>\s*\d+\.\d+\.\d+\.\d+)\."; // [MJ-1/18/2014]: Seems that the product version can be preceeded by a space... go figure..
        private const string REGEX_NEW_FILE_VERSION = @"INFO: szVersionString\(len=\d+\): (?<FILEVERSION>\d+\.\d+\.\d+\.\d+)\.";
        private const string REGEX_ASSEMBLY_VERSION = @"2\)\s+Assembly Version\s+: (?<ASSEMBLYVERSION>\d+\.\d+\.\d+\.\d+)\.";
        private const string REGEX_NEW_ASSEMBLY_VERSION = @"INFO: szAssemblyVersionString\(len=\d+\): (?<ASSEMBLYVERSION>\d+\.\d+\.\d+\.\d+)\.";
        private const string REGEX_SERIALIZER_PARENTASSEMBLYID = @"Info:\s+Serializer-ParentAssemblyId:\s+(?<PARENTASSEMBLYID>(\w|[-,])+)\.";
        private const string REGEX_SERIALIZER_VERSION = @"Info:\s+Serializer-Version:\s+(?<VERSION>\d+\.\d+\.\d+\.\d+)\.";
        private const string REGEX_NEW_SERIALIZER_PARENTASSEMBLYID = @"Info:\s+NewSerializerParentAssemblyId (ANSI):\s+(?<PARENTASSEMBLYID>(\w|[-,])+)\.";
        private const string REGEX_NEW_SERIALIZER_VERSION = @"Info:\s+NewSerializerVersion (ANSI):\s+(?<VERSION>\d+\.\d+\.\d+\.\d+)\.";
        private const string REGEX_PDB_FILE_PATH_IN_DEBUGINFO = @"PDB Path:\s+(?<PDBPATH>[^,]+),";
        private const string REGEX_CHECKSUM_IN_OPTIONAL_HEADER = @"Checksum \(from optional header\):\s+(?<CHECKSUM>\d+)\s+\([^.]+\.";
        private const string REGEX_PRIV_IMPL_DETAILS_STRING = @"PrivateImplementationDetails: <PrivateImplementationDetails>(?<PRIVIMPLDETAILSGUID>{[^}]+}):\d+:\d+\.";
        private const string GET_HEADER_TEMPLATE = "\"{0}\"";
        private const string SET_HEADER_TEMPLATE = "\"{0}\" /c {1} {2} /t {3} /v \"{4}\" /a \"{5}\" /g {6} {7} /d \"{8}\" /k {9} /p \"{10}\"{11}";
        private const string SET_HEADER_WITHOUT_VERSION_TEMPLATE = "\"{0}\" /c {1} {2} /t {3} {4} {5} /g {6} {7} /d \"{8}\" /k {9} /p \"{10}\"{11}";
        private const string FILE_VERSION_HEADER_TEMPLATE = "/v \"{0}\"";
        private const string ASSEMBLY_VERSION_HEADER_TEMPLATE = "/a \"{0}\"";
        private const string PRIVATE_IMPLEMENTATION_DETAILS_HEADER_TEMPLATE = " /r {0}";

        #endregion

        #region Utility Methods

        public static SimpleHeader GetSimpleHeader(string imageFilePath, string toolPath, bool showTraceInfo)
        {
            try
            {
                string arguments = string.Format(GET_HEADER_TEMPLATE, imageFilePath);
                string commandOutput = executeCommand(toolPath, arguments, showTraceInfo);

                if (string.IsNullOrWhiteSpace(commandOutput))
                    return null;

                return parseCommandOutputForCurrentValues(commandOutput);
            }
            catch (Exception ex)
            {
                Trace.TraceError(ex.Message);
                Trace.TraceError(ex.StackTrace);

                return null;
            }
        }

        public static SimpleHeader SetSimpleHeader(string fileHeaderMvid, string debugDirectoryEntryMvid, string timeDateStamp, string fileVersion, string assemblyFileVersion, string serializerParentAssemblyId, string serializerVersion, string imageFilePath, string toolPath, bool showTraceInfo)
        {
            return SetSimpleHeader(new SimpleHeader(fileHeaderMvid, debugDirectoryEntryMvid, timeDateStamp, fileVersion, assemblyFileVersion, serializerParentAssemblyId, serializerVersion), imageFilePath, toolPath, showTraceInfo);
        }

        public static SimpleHeader SetSimpleHeader(SimpleHeader simpleHeader, string imageFilePath, string toolPath, bool showTraceInfo)
        {
            try
            {
                string arguments = string.Format(SET_HEADER_TEMPLATE,
                    imageFilePath,
                    simpleHeader.FileHeaderMvid,
                    simpleHeader.DebugDirectoryEntryMvid,
                    simpleHeader.TimeDateStamp,
                    simpleHeader.FileVersion,
                    simpleHeader.AssemblyVersion,
                    simpleHeader.SerializerParentAssemblyId,
                    simpleHeader.SerializerVersion,
                    simpleHeader.PdbPath,
                    simpleHeader.Checksum,
                    simpleHeader.ProductVersion ?? string.Empty,
                    simpleHeader.PrivateImplementationDetails.Length == 0 ? string.Empty : string.Format(PRIVATE_IMPLEMENTATION_DETAILS_HEADER_TEMPLATE, string.Join(" ", simpleHeader.PrivateImplementationDetails)));

                string commandOutput = executeCommand(toolPath, arguments, showTraceInfo);

                if (string.IsNullOrWhiteSpace(commandOutput))
                    return null;

                return parseCommandOutputForNewValues(commandOutput);
            }
            catch (Exception ex)
            {
                Trace.TraceError(ex.Message);
                Trace.TraceError(ex.StackTrace);

                return null;
            }
        }

        public static SimpleHeader SetSimpleHeaderWithoutVersionChange(SimpleHeader simpleHeader, string imageFilePath, string toolPath, bool keepFileVersion, bool keepAssemblyVersion, bool showTraceInfo)
        {
            try
            {
                string arguments = string.Format(SET_HEADER_WITHOUT_VERSION_TEMPLATE,
                    imageFilePath,
                    simpleHeader.FileHeaderMvid,
                    simpleHeader.DebugDirectoryEntryMvid,
                    simpleHeader.TimeDateStamp,
                    string.Format(FILE_VERSION_HEADER_TEMPLATE, keepFileVersion ? string.Empty : simpleHeader.FileVersion),
                    keepAssemblyVersion ? string.Empty : string.Format(ASSEMBLY_VERSION_HEADER_TEMPLATE, simpleHeader.AssemblyVersion),
                    simpleHeader.SerializerParentAssemblyId,
                    simpleHeader.SerializerVersion,
                    simpleHeader.PdbPath,
                    simpleHeader.Checksum,
                    simpleHeader.ProductVersion ?? string.Empty,
                    simpleHeader.PrivateImplementationDetails.Length == 0 ? string.Empty : string.Format(PRIVATE_IMPLEMENTATION_DETAILS_HEADER_TEMPLATE, string.Join(" ", simpleHeader.PrivateImplementationDetails)));

                string commandOutput = executeCommand(toolPath, arguments, showTraceInfo);

                if (string.IsNullOrWhiteSpace(commandOutput))
                    return null;

                return parseCommandOutputForNewValues(commandOutput);
            }
            catch (Exception ex)
            {
                Trace.TraceError(ex.Message);
                Trace.TraceError(ex.StackTrace);

                return null;
            }
        }

        #endregion

        #region Properties

        public static string DefaultToolPath
        {
            get
            {
                return PE_HEADER_TOOLPATH;
            }
        }

        #endregion

        #region Private Helper Methods

        private static string executeCommand(string toolPath, string arguments, bool showTraceInfo)
        {
            int exitCode = 0;
            try
            {
                ProcessStartInfo processInfo = new ProcessStartInfo();
                processInfo.CreateNoWindow = true;
                processInfo.UseShellExecute = false;
                processInfo.WindowStyle = ProcessWindowStyle.Hidden;
                //processInfo.WorkingDirectory = "";
                processInfo.RedirectStandardError = true;
                processInfo.RedirectStandardOutput = true;
                processInfo.FileName = toolPath;
                Trace.TraceInformation(MESSAGE_COMMAND_CURRENTLY_EXECUTED, PE_HEADER_TOOLPATH, arguments);
                processInfo.Arguments = arguments;
                using (Process process = new Process())
                {
                    process.StartInfo = processInfo;

                    ProcessOutput processOutput = new ProcessOutput();

                    process.OutputDataReceived += processOutput.OutputDataReceived;
                    process.ErrorDataReceived += processOutput.ErrorDataReceived;

                    bool result = process.Start();

                    process.BeginOutputReadLine();
                    process.BeginErrorReadLine();

                    process.WaitForExit(MAXIMUM_PROCESS_EXIT_WAIT_MILLISECONDS);

                    // Read what is already there in the standard output and standard error streams.
                    string outputString = processOutput.StandardOutput;
                    string errorString = processOutput.StandardError;

                    if (showTraceInfo)
                    {
                        Trace.WriteLine(outputString);
                        Trace.WriteLine(errorString);
                    }

                    if (process.HasExited)
                        exitCode = process.ExitCode;
                    else
                    {
                        process.Kill();
                        Trace.TraceError(
                            MESSAGE_FAILED_COMMAND_TIMEOUT,
                            MAXIMUM_PROCESS_EXIT_WAIT_MILLISECONDS);
                        exitCode = FAILURE_EXIT_CODE;
                    }

                    if (!result || exitCode != 0)
                    {
                        Trace.TraceError(
                            MESSAGE_FAILED_RUNNING_COMMAND_PROCESS,
                            result.ToString(), exitCode, PE_HEADER_TOOLPATH, arguments);
                        exitCode = FAILURE_EXIT_CODE;

                        return null;
                    }

                    return outputString;
                }
            }
            catch (Exception ex)
            {
                Trace.TraceError(ex.Message);
                Trace.TraceError(ex.StackTrace);

                return null;
            }
        }

        private static SimpleHeader parseCommandOutputForCurrentValues(string commandOutput)
        {
            try
            {
                SimpleHeader simpleHeader = new SimpleHeader();

                Match fileVersionMatch = Regex.Match(commandOutput, REGEX_FILE_VERSION, RegexOptions.IgnoreCase | RegexOptions.CultureInvariant);
                if (fileVersionMatch != null && fileVersionMatch.Groups != null && fileVersionMatch.Groups.Count > 0)
                    simpleHeader.FileVersion = fileVersionMatch.Groups["FILEVERSION"].Value;

                Match assemblyFileVersionMatch = Regex.Match(commandOutput, REGEX_ASSEMBLY_VERSION, RegexOptions.IgnoreCase | RegexOptions.CultureInvariant);
                if (assemblyFileVersionMatch != null && assemblyFileVersionMatch.Groups != null && assemblyFileVersionMatch.Groups.Count > 0)
                    simpleHeader.AssemblyVersion = assemblyFileVersionMatch.Groups["ASSEMBLYVERSION"].Value;

                Match timeDateStampMatch = Regex.Match(commandOutput, REGEX_TIMEDATESTAMP, RegexOptions.IgnoreCase | RegexOptions.CultureInvariant);
                if (timeDateStampMatch != null && timeDateStampMatch.Groups != null && timeDateStampMatch.Groups.Count > 0)
                    simpleHeader.TimeDateStamp = timeDateStampMatch.Groups["TIMEDATESTAMP"].Value;

                Match fileHeaderMvidMatch = Regex.Match(commandOutput, REGEX_MVID_FILE_HEADER, RegexOptions.IgnoreCase | RegexOptions.CultureInvariant);
                if (fileHeaderMvidMatch != null && fileHeaderMvidMatch.Groups != null && fileHeaderMvidMatch.Groups.Count > 0)
                    simpleHeader.FileHeaderMvid = fileHeaderMvidMatch.Groups["MVID_FILE_HEADER"].Value;

                Match debugDirectoryEntryMvidMatch = Regex.Match(commandOutput, REGEX_MVID_DEBUG_DIRECTORY_ENTRY, RegexOptions.IgnoreCase | RegexOptions.CultureInvariant);
                if (debugDirectoryEntryMvidMatch != null && debugDirectoryEntryMvidMatch.Groups != null && debugDirectoryEntryMvidMatch.Groups.Count > 0)
                    simpleHeader.DebugDirectoryEntryMvid = debugDirectoryEntryMvidMatch.Groups["MVID_FILE_HEADER"].Value;

                Match serializerParentAssemblyIdMatch = Regex.Match(commandOutput, REGEX_SERIALIZER_PARENTASSEMBLYID, RegexOptions.IgnoreCase | RegexOptions.CultureInvariant);
                if (serializerParentAssemblyIdMatch != null && serializerParentAssemblyIdMatch.Groups != null && serializerParentAssemblyIdMatch.Groups.Count > 0)
                    simpleHeader.SerializerParentAssemblyId = serializerParentAssemblyIdMatch.Groups["PARENTASSEMBLYID"].Value;

                Match serializerVersionMatch = Regex.Match(commandOutput, REGEX_SERIALIZER_VERSION, RegexOptions.IgnoreCase | RegexOptions.CultureInvariant);
                if (serializerVersionMatch != null && serializerVersionMatch.Groups != null && serializerVersionMatch.Groups.Count > 0)
                    simpleHeader.SerializerVersion = serializerVersionMatch.Groups["VERSION"].Value;

                Match pdbPathMatch = Regex.Match(commandOutput, REGEX_PDB_FILE_PATH_IN_DEBUGINFO, RegexOptions.IgnoreCase | RegexOptions.CultureInvariant);
                if (pdbPathMatch != null && pdbPathMatch.Groups != null && pdbPathMatch.Groups.Count > 0)
                    simpleHeader.PdbPath = pdbPathMatch.Groups["PDBPATH"].Value;

                Match checksumMatch = Regex.Match(commandOutput, REGEX_CHECKSUM_IN_OPTIONAL_HEADER, RegexOptions.IgnoreCase | RegexOptions.CultureInvariant);
                if (checksumMatch != null && checksumMatch.Groups != null && checksumMatch.Groups.Count > 0)
                    simpleHeader.Checksum = checksumMatch.Groups["CHECKSUM"].Value;

                Match productVersionMatch = Regex.Match(commandOutput, REGEX_PRODUCT_VERSION, RegexOptions.IgnoreCase | RegexOptions.CultureInvariant);
                if (productVersionMatch != null && productVersionMatch.Groups != null && productVersionMatch.Groups.Count > 0)
                    simpleHeader.ProductVersion = productVersionMatch.Groups["PRODUCTVERSION"].Value;

                MatchCollection privateImplDetailsMatches = Regex.Matches(commandOutput, REGEX_PRIV_IMPL_DETAILS_STRING, RegexOptions.IgnoreCase | RegexOptions.CultureInvariant);
                if (privateImplDetailsMatches.Count > 0)
                {
                    List<string> privateImplDetailsList = new List<string>();

                    foreach (Match privateImplDetailsMatch in privateImplDetailsMatches)
                    {
                        if (privateImplDetailsMatch.Groups.Count > 0)
                        {
                            privateImplDetailsList.Add(privateImplDetailsMatch.Groups["PRIVIMPLDETAILSGUID"].Value);
                        }
                    }

                    simpleHeader.PrivateImplementationDetails = privateImplDetailsList.ToArray();
                }

                return simpleHeader;
            }
            catch (Exception ex)
            {
                Trace.TraceError(ex.Message);
                Trace.TraceError(ex.StackTrace);

                return null;
            }
        }

        private static SimpleHeader parseCommandOutputForNewValues(string commandOutput)
        {
            try
            {
                SimpleHeader simpleHeader = new SimpleHeader();

                Match fileVersionMatch = Regex.Match(commandOutput, REGEX_NEW_FILE_VERSION, RegexOptions.IgnoreCase | RegexOptions.CultureInvariant);
                if (fileVersionMatch != null && fileVersionMatch.Groups != null && fileVersionMatch.Groups.Count > 0)
                    simpleHeader.FileVersion = fileVersionMatch.Groups["FILEVERSION"].Value;

                Match assemblyFileVersionMatch = Regex.Match(commandOutput, REGEX_NEW_ASSEMBLY_VERSION, RegexOptions.IgnoreCase | RegexOptions.CultureInvariant);
                if (assemblyFileVersionMatch != null && assemblyFileVersionMatch.Groups != null && assemblyFileVersionMatch.Groups.Count > 0)
                    simpleHeader.AssemblyVersion = assemblyFileVersionMatch.Groups["ASSEMBLYVERSION"].Value;

                Match timeDateStampMatch = Regex.Match(commandOutput, REGEX_NEW_TIMEDATESTAMP, RegexOptions.IgnoreCase | RegexOptions.CultureInvariant);
                if (timeDateStampMatch != null && timeDateStampMatch.Groups != null && timeDateStampMatch.Groups.Count > 0)
                    simpleHeader.TimeDateStamp = timeDateStampMatch.Groups["TIMEDATESTAMP"].Value;

                Match fileHeaderMvidMatch = Regex.Match(commandOutput, REGEX_NEW_MVID_FILE_HEADER, RegexOptions.IgnoreCase | RegexOptions.CultureInvariant);
                if (fileHeaderMvidMatch != null && fileHeaderMvidMatch.Groups != null && fileHeaderMvidMatch.Groups.Count > 0)
                    simpleHeader.FileHeaderMvid = fileHeaderMvidMatch.Groups["MVID_FILE_HEADER"].Value;

                Match debugDirectoryEntryMvidMatch = Regex.Match(commandOutput, REGEX_NEW_MVID_DEBUG_DIRECTORY_ENTRY, RegexOptions.IgnoreCase | RegexOptions.CultureInvariant);
                if (debugDirectoryEntryMvidMatch != null && debugDirectoryEntryMvidMatch.Groups != null && debugDirectoryEntryMvidMatch.Groups.Count > 0)
                    simpleHeader.DebugDirectoryEntryMvid = debugDirectoryEntryMvidMatch.Groups["MVID_FILE_HEADER"].Value;

                Match serializerParentAssemblyIdMatch = Regex.Match(commandOutput, REGEX_NEW_SERIALIZER_PARENTASSEMBLYID, RegexOptions.IgnoreCase | RegexOptions.CultureInvariant);
                if (serializerParentAssemblyIdMatch != null && serializerParentAssemblyIdMatch.Groups != null && serializerParentAssemblyIdMatch.Groups.Count > 0)
                    simpleHeader.SerializerParentAssemblyId = serializerParentAssemblyIdMatch.Groups["PARENTASSEMBLYID"].Value;

                Match serializerVersionMatch = Regex.Match(commandOutput, REGEX_NEW_SERIALIZER_VERSION, RegexOptions.IgnoreCase | RegexOptions.CultureInvariant);
                if (serializerVersionMatch != null && serializerVersionMatch.Groups != null && serializerVersionMatch.Groups.Count > 0)
                    simpleHeader.SerializerVersion = serializerVersionMatch.Groups["VERSION"].Value;

                Match pdbPathMatch = Regex.Match(commandOutput, REGEX_PDB_FILE_PATH_IN_DEBUGINFO, RegexOptions.IgnoreCase | RegexOptions.CultureInvariant);
                if (pdbPathMatch != null && pdbPathMatch.Groups != null && pdbPathMatch.Groups.Count > 0)
                    simpleHeader.PdbPath = pdbPathMatch.Groups["PDBPATH"].Value;

                Match checksumMatch = Regex.Match(commandOutput, REGEX_CHECKSUM_IN_OPTIONAL_HEADER, RegexOptions.IgnoreCase | RegexOptions.CultureInvariant);
                if (checksumMatch != null && checksumMatch.Groups != null && checksumMatch.Groups.Count > 0)
                    simpleHeader.Checksum = checksumMatch.Groups["CHECKSUM"].Value;

                Match productVersionMatch = Regex.Match(commandOutput, REGEX_PRODUCT_VERSION, RegexOptions.IgnoreCase | RegexOptions.CultureInvariant);
                if (productVersionMatch != null && productVersionMatch.Groups != null && productVersionMatch.Groups.Count > 0)
                    simpleHeader.ProductVersion = productVersionMatch.Groups["PRODUCTVERSION"].Value;

                MatchCollection privateImplDetailsMatches = Regex.Matches(commandOutput, REGEX_PRIV_IMPL_DETAILS_STRING, RegexOptions.IgnoreCase | RegexOptions.CultureInvariant);
                if (privateImplDetailsMatches.Count > 0)
                {
                    List<string> privateImplDetailsList = new List<string>();

                    foreach (Match privateImplDetailsMatch in privateImplDetailsMatches)
                    {
                        if (privateImplDetailsMatch.Groups.Count > 0)
                        {
                            privateImplDetailsList.Add(privateImplDetailsMatch.Groups["PRIVIMPLDETAILSGUID"].Value);
                        }
                    }

                    simpleHeader.PrivateImplementationDetails = privateImplDetailsList.ToArray();
                }

                return simpleHeader;
            }
            catch (Exception ex)
            {
                Trace.TraceError(ex.Message);
                Trace.TraceError(ex.StackTrace);

                return null;
            }
        }

        #endregion

    }
}
